#!/usr/bin/env python

# Copyright 2012 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Street Identity AP Demo Application."""

import cgi

try:  # pragma: no cover
  import simplejson
except ImportError:  # pragma: no cover
  try:
    # Try to import from django, should work on App Engine
    from django.utils import simplejson
  except ImportError:
    # Should work for Python2.6 and higher.
    import json as simplejson
import logging
import os
import pickle
import random
import common
from common import SUCCESSFUL_LOGIN_REDIRECT
from google.appengine.api import memcache
from google.appengine.ext import db
from google.appengine.ext import webapp
from google.appengine.ext.webapp import template
from google.appengine.ext.webapp.util import run_wsgi_app
from lib import authn
import model
from oauth2client.client import FlowExchangeError
from openidconnect.error import FlowTokenInfoError
from openidconnect.error import FlowUserInfoError
from streetidentity_ap import oauth
from streetidentity_ap.client import StreetIdentityAPFlow
from streetidentity_ap.error import RegistrationError
import utils

__author__ = 'Maciej Machulak'
__maintainer__ = 'Maciej Machulak'
__email__ = 'mmachulak@google.com'

__copyright__ = 'Copyright 2012, Google Inc.'
__license__ = 'Apache License 2.0'
__version__ = '0.2'
__status__ = 'Prototype'

SESSION_ID = common.SESSION_ID
CALLBACK_URL = '/oauth2callback'
CREDENTIALS_PREFIX = 'credentials::'
ATTRIBUTES_PREFIX = 'attributes::'
ADDRESS_TYPE = 'street_address'
PHONE_TYPE = 'phone_number'
DOB_TYPE = 'birthday'
HEADER_SET_COOKIE = 'Set-Cookie'

_CLIENTID = '776695651878.apps.googleusercontent.com'

IDP_FLOW = StreetIdentityAPFlow(
    client_id=_CLIENTID,
    client_secret='',
    user_agent='Street Identity AP Demo Application',
    attribute_uri='https://streetidentity-console.appspot'
    '.com/api/v1/attribute_info/')


class AuthNError(Exception):
  """Raised when there is an error with the login component."""

  def __init__(self, msg):
    Exception.__init__(self)
    self.msg = msg

  def __str__(self):
    return repr(self.msg)


class MainHandler(webapp.RequestHandler):
  """Handles the main entry to the application."""

  @authn.redirect_logged_user(SUCCESSFUL_LOGIN_REDIRECT)
  def get(self):
    path = os.path.join(os.path.dirname(__file__), 'templates/index.html')
    self.response.out.write(template.render(path, {}))


class SignInHandler(webapp.RequestHandler):
  """Generates a sign in link for the correct identity provider."""

  def get(self):
    callback = self.request.relative_url(CALLBACK_URL)
    logging.debug('Callback: %s', callback)
    authorize_url = IDP_FLOW.step1_get_authorize_url(callback)
    logging.debug('Authorization URL; %s', authorize_url)
    self.redirect(authorize_url)


class SignOutHandler(webapp.RequestHandler):
  """Destroys the session for the user."""

  @authn.login_required()
  def get(self):
    self.response.headers.add_header(HEADER_SET_COOKIE, SESSION_ID + '=')
    memcache.delete(self.session.id)
    self.session.delete()
    path = os.path.join(os.path.dirname(__file__), 'templates/index.html')
    self.response.out.write(template.render(path, {}))


class AccountHandler(webapp.RequestHandler):
  """Account handler for the user - user must be authenticated."""

  @authn.login_required()
  def get(self):
    """Show account page."""

    attributes = model.Attribute.all().filter('user = ', self.user)
    attribute_list = {}
    if attributes:
      for attribute in attributes:
        attribute_list[attribute.type] = attribute

    values = {
        'email': self.user.email,
        'street_address': attribute_list.get('street_address'),
        'phone': attribute_list.get('phone_number'),
        'dob': attribute_list.get('birthday'),
    }
    path = os.path.join(os.path.dirname(__file__), 'templates/account.html')
    self.response.out.write(template.render(path, values))


class OAuthHandler(webapp.RequestHandler):
  """.OAuth Redirect handler."""

  def get(self):
    logging.debug('exchanging code for access token')
    try:
      credentials = IDP_FLOW.step234_exchange_and_tokeninfo_and_userinfo(
          self.request.params)
      user_info = credentials.userinfo
      logging.debug('obtained user info')

      # Required user attributes
      email = user_info.email
      display_name = user_info.name
      refresh_token = credentials.refresh_token

      # Retrieve existing user
      user = model.User.all().filter('email = ', email).get()
      if user is None:
        # If user does not exist yet, create a new user
        logging.debug('user not found - creating a new one')

        user = model.User(email=email,
                          display_name=display_name,
                          refresh_token=refresh_token,
                          credentials=db.Blob(pickle.dumps(credentials)))
        user.put()

      # Update the user
      logging.debug('updating the user in the datastore with new credentials')
      user.credentials = db.Blob(pickle.dumps(credentials))
      user.put()

      # Access token for the user
      access_token = credentials.access_token

      # Generate session id
      key = str(random.random())[2:] + str(random.random())[2:]
      logging.debug('generated new session id: %s ', key)

      # Create session for the user and store it in datastore
      session = model.Session(id=key, user=user, access_token=access_token)
      session.put()
      logging.debug('session stored in the datastore')

      # Store the session in memcache
      memcache.add(session.id, session)
      logging.debug('session stored in memcache')

      # Set the cookie
      self.response.headers.add_header(HEADER_SET_COOKIE,
                                       SESSION_ID + '=' + key + '; Path=/')
      logging.debug('set cookie with session id')

      self.redirect('/account')
    except FlowExchangeError or FlowTokenInfoError or FlowUserInfoError:
      logging.error('error during OpenID Connect flow')
      values = {
          'error_message': 'could not complete OpenID Connect flow'
      }
      path = os.path.join(os.path.dirname(__file__), 'templates/error.html')
      self.response.out.write(template.render(path, values))
      return


class AttributeHandler(webapp.RequestHandler):
  """Handles attribute registration by the End-User."""

  @authn.login_required()
  def post(self):
    """Register a new attribute."""

    logging.debug('creating a new attribute')
    type_ = cgi.escape(self.request.get('type'))
    value = cgi.escape(self.request.get('value'))
    desc = cgi.escape(self.request.get('desc'))

    attribute = model.Attribute(id=utils.random(),
                                type=type_,
                                value=value,
                                desc=desc,
                                user=self.user)
    attribute.put()
    logging.debug('attribute stored in datastore')

    self.redirect('/account')


class StreetIdentitySaveHandler(webapp.RequestHandler):
  """Handles a user's request to register attribute with IDP."""

  @authn.login_required()
  def post(self):
    """Save attribute metadata to IDP."""
    id_ = cgi.escape(self.request.get('id'))
    logging.debug('received request to save attribute with id: %s', id_)

    attribute = model.Attribute.all().filter('id = ', id_).get()

    if attribute is None:
      logging.error('could not find attribute')
      self.redirect('/account')
      return

    logging.debug('attribute retrieved from datastore')

    # save to google
    try:
      credentials = pickle.loads(str(self.user.credentials))
      registration_response = IDP_FLOW.register_attribute(attribute.type,
                                                          attribute.desc,
                                                          credentials)
      logging.debug('attribute has been successfully registered at IDP')

      attribute.si_id = registration_response.id
      attribute.si_location = registration_response.location
      attribute.put()

      logging.debug('attribute stored within the datastore')

      self.redirect('/account')
      return

    except RegistrationError:
      logging.error('could not register attribute at IDP')
      self.redirect('/account')
      return


class StreetIdentityRemoveHandler(webapp.RequestHandler):
  """Handles a user's request to deregister attribute from IDP."""

  @authn.login_required()
  def post(self):
    """Save attribute metadata to IDP."""

    id_ = cgi.escape(self.request.get('id'))
    logging.debug('received request to remove attribute with id: %s', id_)

    attribute = model.Attribute.all().filter('id = ', id_).get()
    if attribute is None:
      logging.error('could not find attribute')
      self.redirect('/account')
      return

    logging.debug('attribute retrieved from datastore')

    # remove attribute from google
    try:
      credentials = pickle.loads(str(self.user.credentials))
      IDP_FLOW.remove_attribute(attribute.si_location, credentials)
      logging.debug('attribute has been successfully removed from the IDP')

      attribute.si_id = ''
      attribute.si_location = ''
      attribute.put()

      self.redirect('/account')
      return

    except RegistrationError:
      logging.error('could not remove the attribute from the IDP')
      self.redirect('/account')
      return


class StreetIdentityAPIHandler(webapp.RequestHandler):
  """Street Identity API at Attribute Provider."""

  @oauth.api_authz_required()
  def get(self, path):
    """Return attribute."""

    logging.debug('attribute API')

    # Get token from the request
    jwt_token = self.request.get('si_access_token', None)
    if not jwt_token:
      logging.error('JWT token missing - return with error')
      self.response.set_status(403)
      return

    # Load credentials
    credentials = pickle.loads(str(self.user.credentials))

    # Check authorization for this attribute
    verified_token = IDP_FLOW.verify_token(token=jwt_token, audience=_CLIENTID,
                                           credentials=credentials)

    # Get the id of the resource
    id_ = path
    logging.debug('id of the attribute: %s', id_)

    # Check if the scope is correct
    if verified_token.scope != self.request.relative_url(id_):
      logging.error('JWT token valid but issued for a different attribute')
      self.response.set_status(403)
      return

    # Return the attribute
    logging.debug('retrieving attribute from db with id: %s', id_)
    attribute = model.Attribute.all().filter('si_id = ', id_).get()
    if attribute is None:
      logging.debug('attribute not found - 404')
      self.response.set_status(404)
      return
    logging.debug('attribute found - returning')
    attribute_response = {'type': attribute.type,
                          'desc': attribute.desc,
                          'value': attribute.value}

    self.response.set_status(200)
    self.response.out.write(simplejson.dumps(attribute_response))


class Setup(webapp.RequestHandler):

  def get(self):
    model.Application(client_id='someclient', client_secret='somesecret').put()

routes = [('/', MainHandler),
          ('/sign_in', SignInHandler),
          ('/sign_out', SignOutHandler),
          ('/account', AccountHandler),
          ('/account/add', AttributeHandler),
          ('/streetidentity/save', StreetIdentitySaveHandler),
          ('/streetidentity/remove', StreetIdentityRemoveHandler),
          ('/api/streetidentity/(.*)', StreetIdentityAPIHandler),
          ('/oauth/token', oauth.TokenHandler),
          ('/setup', Setup),
          (CALLBACK_URL, OAuthHandler)]

application = webapp.WSGIApplication(routes, debug=True)


def main():
  logging.getLogger().setLevel(logging.DEBUG)
  run_wsgi_app(application)

if __name__ == '__main__':
  main()
